import os
from playwright.sync_api import sync_playwright
import pyperclip

test = not True
answer_index = 0


_xp = None
p = None
browser = None
page = None
el_input = None
el_send = None
# pwd = 'D:/eglwa/home/apnews'
pwd = os.path.dirname(os.path.abspath(__file__))
personal_data_path = r'C:\Users\eglwa\AppData\Local\Microsoft\Edge\User Data Work'
chat_url = 'https://chat.openai.com/'
input_xpath = '//*[@id="prompt-textarea"]'
send_button_xpath = '//*[@id="__next"]//main//form/div/div[2]/div/button'
copy_content_xpath = '//*[@id="__next"]/div[1]/div[2]/main/div[1]/div[1]/div/div/div/div[{}]/div/div/div[2]/div/div[2]/div/button'
continue_generate_xpath = '//*[@id="__next"]/div[1]/div[2]/div/main/div[2]/form/div/div[1]/div/button[2]'
ua = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36 Edg/114.0.1823.67'

def start():
    global _xp, p, browser, page, el_input, el_send
    _xp = sync_playwright()
    p = _xp.__enter__()
    if test:
        browser = p.chromium.launch_persistent_context(
            channel='msedge', user_data_dir=personal_data_path, headless=False)
    else:
        browser = p.chromium.launch_persistent_context(
            channel='msedge', user_data_dir=personal_data_path, headless=True, user_agent=ua)

    page = browser.new_page()
    page.add_init_script(
        'Object.defineProperty(navigator, "webdriver", {get: () => undefined})')
    page.goto(chat_url)
    browser.set_default_timeout(30*1000)
    el_input = page.query_selector(input_xpath)
    if not el_input:
        raise Exception('Failed to open page')
    el_input = page.locator(input_xpath)
    el_send = page.locator(send_button_xpath)
    browser.set_default_timeout(60*1000)


def screenshot():
    page.screenshot(path=f'{pwd}/screenshot.png')


def make_a_question(question):
    global answer_index
    if not page.query_selector(input_xpath):
        raise Exception('unexpected Error')
    if test:
        page.screenshot(path=f'{pwd}/screenshot.png')
    bxs = page.query_selector_all('//*[@id="__next"]//main/div[1]/div[1]/div/div/div')
    count = 0
    if bxs:
        if len(bxs) == 2 and bxs[0].text_content()=='GPT-3.5GPT-4':
            count = 0
        else:
            count = len(bxs)
            if count % 2 == 1:
                count -= 1
    else:
        count = answer_index
    el_input.fill(question)
    el_send.click()
    answer_index = count+2 
    copy_content_xpath_ = copy_content_xpath.format(answer_index)
    el_copy = page.locator(copy_content_xpath_)
    el_copy.click()
    continue_el = page.query_selector(continue_generate_xpath)
    if continue_el:
        continue_el.click()
        el_copy = page.locator(copy_content_xpath)
    el_copy.click()
    if not test:
        content_xpath = f'//*[@id="__next"]//main/div/div[1]/div/div/div/div[{answer_index}]'
        el_content = page.locator(content_xpath)
        answer = el_content.text_content()
        if answer.startswith('ChatGPTChatGPT1 / 1'):
            answer = answer[19:]
    else:
        answer = pyperclip.paste()
    if not page.query_selector(input_xpath):
        raise Exception('unexpected Error')
    if test:
        page.screenshot(path=f'{pwd}/screenshot.png')

    return answer


def new_topic():
    global answer_index
    close(False)
    new_topic_btn_xpath = '//*[@id="__next"]/div/div/div[1]/div[1]/div/div/div/nav/div[1]/a'
    page.locator(new_topic_btn_xpath).click()
    answer_index = 0


def close(exit=True):
    # delete session
    session_xpath = '//*[@id="__next"]/div/div/div[1]/div[1]/div/div/div/nav/div[3]/div/div/span[1]/div[1]/ol/li[1]'
    delete_session_xpath = '//*[@id="__next"]/div/div/div[1]/div[1]/div/div/div/nav/div[3]/div/div/span[1]/div[1]/ol/li[1]/a/div[2]/button[2]'
    confirm_xpath = 'button.btn.relative.btn-danger' 
    page.locator(session_xpath).click()
    page.locator(delete_session_xpath).click()
    page.locator(confirm_xpath).click()
    if not exit:
        return
    quit()


def quit():
    browser.close()
    _xp.__exit__(None, None, None)


def reload():
    global _xp, p, browser, page, el_input, el_send, answer_index
    page.reload()
    el_input = page.locator(input_xpath)
    el_send = page.locator(send_button_xpath)
    answer_index = 0


if __name__ == '__main__':
    start()
    while 1:
        x = input('>>')
        if x == 'q':
            break
        if not x:
            continue
        answer = make_a_question(x)
        print(answer)
    close()
